package longshu.data.common.ds;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.jdbc.datasource.lookup.DataSourceLookup;
import org.springframework.jdbc.datasource.lookup.DataSourceLookupFailureException;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * DynamicDataSource 动态数据源(单例)
 *
 * @author LongShu 2018/01/18
 * @see org.springframework.jdbc.datasource.lookup.MapDataSourceLookup
 */
public class DynamicDataSource extends AbstractRoutingDataSource implements DataSourceLookup {

    private final ThreadLocal<String> dataSourceKey = new ThreadLocal<>();
    private final Map<String, DataSource> dataSources = new HashMap<>(4);

    private static final DynamicDataSource INSTANCE = new DynamicDataSource();

    private DynamicDataSource() {
    }

    public static DynamicDataSource getInstance() {
        return INSTANCE;
    }

    /**
     * @see #clear()
     */
    public static void setDataSourceKey(String dsKey) {
        getInstance().dataSourceKey.set(dsKey);
    }

    public static void clear() {
        getInstance().dataSourceKey.remove();
    }

    public static void addDataSource(String dsKey, DataSource dataSource) {
        getInstance().dataSources.put(dsKey, dataSource);
    }

    /**
     * @return 当前数据源
     * @see #setDataSourceKey(String)
     * @see #determineTargetDataSource()
     */
    @SuppressWarnings("unchecked")
    public static <T extends DataSource> T currentDataSource() {
        DataSource dataSource = getInstance().determineTargetDataSource();
        return (T) dataSource;
    }

    @Override
    public void afterPropertiesSet() {
        /**
         * @see  #resolveSpecifiedDataSource(Object)
         */
        super.setDataSourceLookup(this);
        super.setTargetDataSources(new HashMap<>(dataSources));

        super.afterPropertiesSet();
    }

    @Override
    public Object determineCurrentLookupKey() {
        /**
         * 动态切换
         * @see #determineTargetDataSource()
         */
        return dataSourceKey.get();
    }

    /**
     * @see #resolveSpecifiedDataSource(Object)
     */
    @Override
    public DataSource getDataSource(String dataSourceName) throws DataSourceLookupFailureException {
        DataSource dataSource = dataSources.get(dataSourceName);
        if (dataSource == null) {
            throw new DataSourceLookupFailureException("没有找到:" + dataSourceName);
        }
        return dataSource;
    }

}